from torch.utils.tensorboard import SummaryWriter
import os, utils, glob, losses
import sys
from torch.utils.data import DataLoader
from data import datasets, trans
import numpy as np
import torch
from torchvision import transforms
from torch import optim
import matplotlib.pyplot as plt
from natsort import natsorted
from models.TransMorph_ConvMultiDWin_Separate_Chan_SEF_XMLP_TVF_SPT import CONFIGS as CONFIGS_TM
import models.TransMorph_ConvMultiDWin_Separate_Chan_SEF_XMLP_TVF_SPT as TransMorph
import torch.nn.functional as F
from scipy.ndimage.interpolation import map_coordinates, zoom
from surface_distance import compute_dice_coefficient

class Logger(object):
    def __init__(self, save_dir):
        self.terminal = sys.stdout
        self.log = open(save_dir+"logfile.log", "a")

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)

    def flush(self):
        pass

def main():
    batch_size = 1
    train_dir = 'D:/DATA/OASIS/All/'
    val_dir = 'D:/DATA/OASIS/Test/'
    weights = [1, 1, 1] # loss weights
    Dwin = [7, 5, 3]
    time_steps = 12
    save_dir = 'TransMorphConvMultiDWinSep_SPT_DWin_{}{}{}_half_TVF_{}_ncc_{}_dsc{}_diffusion_{}/'.format(Dwin[0], Dwin[1], Dwin[2], time_steps,
                                                                                                              weights[0], weights[1], weights[2])
    if not os.path.exists('experiments/'+save_dir):
        os.makedirs('experiments/'+save_dir)
    if not os.path.exists('logs/'+save_dir):
        os.makedirs('logs/'+save_dir)
    sys.stdout = Logger('logs/'+save_dir)
    lr = 0.0001
    epoch_start = 0
    max_epoch = 500
    '''
    Initialize model
    '''
    H, W, D = 160, 192, 224
    config = CONFIGS_TM['TransMorph-3-LVL']
    config.img_size = (H//2, W//2, D//2)
    config.dwin_kernel_size = (Dwin[0], Dwin[1], Dwin[2])
    config.window_size = (H // 32, W // 32, D // 32)
    model = TransMorph.TransMorphCascadeAd(config, time_steps)
    model.cuda()
    '''
    Initialize spatial transformation function
    '''
    reg_model = utils.register_model((H, W, D), 'nearest')
    reg_model.cuda()
    spatial_trans = TransMorph.SpatialTransformer((H, W, D)).cuda()

    '''
    Initialize training
    '''
    train_composed = transforms.Compose([trans.NumpyType((np.float32, np.int16)),
                                         ])

    val_composed = transforms.Compose([trans.NumpyType((np.float32, np.int16))])
    train_set = datasets.OASISBrainDataset(glob.glob(train_dir + '*.pkl'), transforms=train_composed)
    val_set = datasets.OASISBrainInferDataset(glob.glob(val_dir + '*.pkl'), transforms=val_composed)
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_set, batch_size=1, shuffle=False, num_workers=4, pin_memory=True, drop_last=True)

    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0, amsgrad=True)
    criterion_ncc = losses.NCC_vxm()
    criterion_dsc = losses.DiceLoss()
    criterion_reg = losses.Grad3d(penalty='l2')
    scaler = torch.cuda.amp.GradScaler()
    best_dsc = 0
    writer = SummaryWriter(log_dir='logs/'+save_dir)
    for epoch in range(epoch_start, max_epoch):
        print('Training Starts')
        '''
        Training
        '''
        loss_all = utils.AverageMeter()
        idx = 0
        for data in train_loader:
            idx += 1
            model.train()
            adjust_learning_rate(optimizer, epoch, max_epoch, lr)
            with torch.no_grad():
                x = data[0].half().cuda()
                y = data[1].half().cuda()
                x_half = F.avg_pool3d(x, 2).half().cuda()
                y_half = F.avg_pool3d(y, 2).half().cuda()
                x_seg = data[2]
                x_seg = F.one_hot(x_seg.long().cuda(), 36).float().squeeze(1).permute(0, 4, 1, 2, 3).half().cuda()#F.avg_pool3d(F.one_hot(x_seg.long().cuda(), 36).float().squeeze(1).permute(0, 4, 1, 2, 3), 2).half().cuda()
                y_seg = data[3]
                y_seg = F.one_hot(y_seg.long().cuda(), 36).float().squeeze(1).permute(0, 4, 1, 2, 3).half().cuda()#F.avg_pool3d(F.one_hot(y_seg.long().cuda(), 36).float().squeeze(1).permute(0, 4, 1, 2, 3), 2).half().cuda()
            optimizer.zero_grad()
            with torch.cuda.amp.autocast():
                flow = model((x_half,y_half))
                flow = F.interpolate(flow.cuda(), scale_factor=2, mode='trilinear', align_corners=False) * 2
                output = spatial_trans(x.half(), flow.half())
                def_seg = spatial_trans(x_seg.half(), flow.half())
                loss_ncc = criterion_ncc(output.half(), y.half()) * weights[0]
                loss_dsc = criterion_dsc(def_seg.half(), y_seg.half()) * weights[1]
                loss_reg = criterion_reg(flow.half(), y.half()) * weights[2]
                loss = loss_ncc + loss_dsc + loss_reg
                loss_all.update(loss.item(), y.numel())
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            del def_seg, loss

            optimizer.zero_grad()
            with torch.cuda.amp.autocast():
                flow = model((y_half, x_half))
                flow = F.interpolate(flow.cuda(), scale_factor=2, mode='trilinear', align_corners=False) * 2
                output = spatial_trans(y.half(), flow.half())
                def_seg = spatial_trans(y_seg.half(), flow.half())
                loss_ncc = criterion_ncc(output.half(), x.half()) * weights[0]
                loss_dsc = criterion_dsc(def_seg.half(), x_seg.half()) * weights[1]
                loss_reg = criterion_reg(flow.half(), x.half()) * weights[2]
                loss = loss_ncc + loss_dsc + loss_reg
                loss_all.update(loss.item(), x.numel())
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            del def_seg
            print('Iter {} of {} loss {:.4f}, Img Sim: {:.6f}, DSC: {:.6f}, Reg: {:.6f}'.format(idx, len(train_loader),
                                                                                                loss.item(),
                                                                                                loss_ncc.item(),
                                                                                                loss_dsc.item(),
                                                                                                loss_reg.item()))

        writer.add_scalar('Loss/train', loss_all.avg, epoch)
        print('Epoch {} loss {:.4f}'.format(epoch, loss_all.avg))
        '''
        Validation
        '''
        eval_dsc = utils.AverageMeter()
        with torch.no_grad():
            for data in val_loader:
                model.eval()
                x = data[0]
                y = data[1]
                x = F.avg_pool3d(x, 2).cuda()
                y = F.avg_pool3d(y, 2).cuda()
                x_seg = data[2].cuda()
                y_seg = data[3].cuda()
                grid_img = mk_grid_img(8, 1, (H, W, D))
                output = model((x, y))
                flow = F.interpolate(output.cuda(), scale_factor=2, mode='trilinear') * 2
                disp_field = flow.cpu().detach().numpy()[0]
                disp_field = np.array([zoom(disp_field[i], 0.5, order=2) for i in range(3)]).astype(np.float16).astype(
                    'float32')
                disp_field = np.array([zoom(disp_field[i], 2, order=2) for i in range(3)])
                fixed = y_seg.cpu().detach().numpy()[0, 0]
                moving = x_seg.cpu().detach().numpy()[0, 0]
                D, H, W = fixed.shape
                identity = np.meshgrid(np.arange(D), np.arange(H), np.arange(W), indexing='ij')
                moving_warped = map_coordinates(moving, identity + disp_field, order=0)
                def_grid = spatial_trans(grid_img.float(), flow.cuda())
                dice = 0
                count = 0
                for i in range(1, 36):
                    if ((fixed == i).sum() == 0) or ((moving == i).sum() == 0):
                        continue
                    dice += compute_dice_coefficient((fixed == i), (moving_warped == i))
                    count += 1
                dice /= count
                eval_dsc.update(dice, x.size(0))
                print(eval_dsc.avg)
        best_dsc = max(eval_dsc.avg, best_dsc)
        save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'best_dsc': best_dsc,
            'optimizer': optimizer.state_dict(),
        }, save_dir='experiments/'+save_dir, filename='dsc{:.4f}.pth.tar'.format(eval_dsc.avg))
        writer.add_scalar('DSC/validate', eval_dsc.avg, epoch)
        plt.switch_backend('agg')
        pred_fig = comput_fig(torch.from_numpy(moving_warped[None, None]))
        grid_fig = comput_fig(def_grid)
        x_fig = comput_fig(x_seg)
        tar_fig = comput_fig(y_seg)
        writer.add_figure('Grid', grid_fig, epoch)
        plt.close(grid_fig)
        writer.add_figure('input', x_fig, epoch)
        plt.close(x_fig)
        writer.add_figure('ground truth', tar_fig, epoch)
        plt.close(tar_fig)
        writer.add_figure('prediction', pred_fig, epoch)
        plt.close(pred_fig)
        loss_all.reset()
        del def_grid, grid_img, output
    writer.close()

def comput_fig(img):
    img = img.detach().cpu().numpy()[0, 0, 48:64, :, :]
    fig = plt.figure(figsize=(12,12), dpi=180)
    for i in range(img.shape[0]):
        plt.subplot(4, 4, i + 1)
        plt.axis('off')
        plt.imshow(img[i, :, :], cmap='gray')
    fig.subplots_adjust(wspace=0, hspace=0)
    return fig

def adjust_learning_rate(optimizer, epoch, MAX_EPOCHES, INIT_LR, power=0.9):
    for param_group in optimizer.param_groups:
        param_group['lr'] = round(INIT_LR * np.power( 1 - (epoch) / MAX_EPOCHES ,power),8)

def mk_grid_img(grid_step, line_thickness=1, grid_sz=(160, 192, 224)):
    grid_img = np.zeros(grid_sz)
    for j in range(0, grid_img.shape[1], grid_step):
        grid_img[:, j+line_thickness-1, :] = 1
    for i in range(0, grid_img.shape[2], grid_step):
        grid_img[:, :, i+line_thickness-1] = 1
    grid_img = grid_img[None, None, ...]
    grid_img = torch.from_numpy(grid_img).cuda()
    return grid_img

def save_checkpoint(state, save_dir='models', filename='checkpoint.pth.tar', max_model_num=4):
    torch.save(state, save_dir+filename)
    model_lists = natsorted(glob.glob(save_dir + '*'))
    while len(model_lists) > max_model_num:
        os.remove(model_lists[0])
        model_lists = natsorted(glob.glob(save_dir + '*'))

if __name__ == '__main__':
    '''
    GPU configuration
    '''
    GPU_iden = 1
    GPU_num = torch.cuda.device_count()
    print('Number of GPU: ' + str(GPU_num))
    for GPU_idx in range(GPU_num):
        GPU_name = torch.cuda.get_device_name(GPU_idx)
        print('     GPU #' + str(GPU_idx) + ': ' + GPU_name)
    torch.cuda.set_device(GPU_iden)
    GPU_avai = torch.cuda.is_available()
    print('Currently using: ' + torch.cuda.get_device_name(GPU_iden))
    print('If the GPU is available? ' + str(GPU_avai))
    torch.manual_seed(0)
    main()